# -*- coding: utf-8 -*-
from typing import List
from textwrap import dedent
from airflow import DAG
from airflow.providers.dingding.operators.dingding import DingdingHook
from airflow.models import Variable
import re

# 红色：#EA3323 => 失败
# 绿色：#377E22 => 成功
# 橙色：#FF990F => 上游失败
# 黄色：#FFD211 => 异常，准备重试
# 粉色：#F5C2CB => 跳过
dag_template = """\
# {% if dag_run.state.upper() == 'FAILED' %}【<font color=#EA3323>{{ dag_run.state.upper() }}</font>】\
{% elif dag_run.state.upper() == 'SUCCESS' %}【<font color=#377E22>{{ dag_run.state.upper() }}</font>】\
{% else %}【{{ dag_run.state.upper() }}】{% endif %}{{ dag_run.dag_id }} ({{ var.value.env.upper() }} 环境)

- *任务状态:* {{ dag_run.state.upper() }}
- *任务日期:* {{ execution_date | cst_ds }}
- *运行时长:* {{ dag_run.end_date - dag_run.start_date }}
- *开始时间:* {{ dag_run.start_date | to_cst }}
- *结束结束:* {{ dag_run.end_date | to_cst }}
- *成功任务:* {{ dag_run.get_task_instances(state=['success']) | length }} 
- *跳过任务:* {{ dag_run.get_task_instances(state=['skipped']) | length }} 
- *失败任务:* {{ dag_run.get_task_instances(state=['failed', 'upstream_failed']) | length }} 
{% for ti in dag_run.get_task_instances(state=['failed']) %}\
  > <font color=#EA3323>{{ ti.task_id }}</font>
{% endfor %}
{% for ti in dag_run.get_task_instances(state=['upstream_failed']) %}\
  > <font color=#FF990F>{{ ti.task_id }}(上游失败)</font>
{% endfor %}
"""


def dingding_dag_alert(dingding_conn_id: str = 'dingding_default', at_mobiles: List[str] = None, at_all: bool = False):
    def _dingding_dag_alert(context, *args, **kwargs):
        jinja_env = context['dag'].get_template_env()
        title = jinja_env.from_string("【{{ dag_run.state.upper() }}】{{ dag_run.dag_id }} ({{ var.value.env.upper() }} 环境)").render(context)
        content = dedent(jinja_env.from_string(dag_template).render(context))

        return DingdingHook(dingding_conn_id=dingding_conn_id,
                            message_type='markdown',
                            message={'title': title, 'text': content},
                            at_mobiles=at_mobiles,
                            at_all=at_all, ).send()

    return _dingding_dag_alert


task_template = """
# 任务 {{ ti.task_id }} {% if ti.state.upper() == 'FAILED' %}<font color=#EA3323>{{ ti.state.upper() }}</font>\
{% elif ti.state.upper() == 'SUCCESS' %}<font color=#377E22>{{ ti.state.upper() }}</font>\
{% elif ti.state.upper() == 'UP_FOR_RETRY' %}<font color=#FFD211>{{ ti.state.upper() }}</font>\
{% elif ti.state.upper() == 'SKIPPED' %}<font color=#F5C2CB>{{ ti.state.upper() }}</font>\
{% else %}{{ ti.state.upper() }}{% endif %} ({{ var.value.env.upper() }} 环境)

- *任务日期:* {{ execution_date | cst_ds }}
- *开始时间:* {{ ti.start_date | to_cst }}
"""


def dingding_task_alert(dingding_conn_id: str = 'dingding_default', at_mobiles: List[str] = None, at_all: bool = False):
    def _dingding_task_alert(context, *args, **kwargs):
        jinja_env = context['dag'].get_template_env()
        title = jinja_env.from_string("任务 {{ ti.task_id }} {{ ti.state.upper() }} ({{ var.value.env.upper() }} 环境)").render(context)
        content = dedent(jinja_env.from_string(task_template).render(context))
        contentnew = re.sub(r'"password"((?!"password").)*"driver"', '"password":"******",\n"driver"', content, flags=re.DOTALL)
        print(f'send to dingding: {dingding_conn_id}')
        print(f'title: \n{title}')
        print(f'content: \n{contentnew}')

        return DingdingHook(dingding_conn_id=dingding_conn_id,
                            message_type='markdown',
                            message={'title': title, 'text':contentnew},
                            at_mobiles=at_mobiles,
                            at_all=at_all, ).send()

    return _dingding_task_alert


def dingding_sla_alert(dingding_conn_id: str = 'dingding_default', at_mobiles: List[str] = None, at_all: bool = False):
    def _dingding_sla_alert(dag: DAG, task_list: str, blocking_task_list: str, slas, blocking_tis):
        env = Variable.get('env', '').upper()
        title = f'任务有超时风险[{env}环境]'
        content = f'# 以下任务有<font color=#FFD211>超时</font>风险 ({ env } 环境)\n\n' + '\n'.join(f'- {sla.task_id}' for sla in slas)

        print(f'send to dingding: {dingding_conn_id}')
        print(f'title: \n{title}')
        print(f'content: \n{content}')

        return DingdingHook(dingding_conn_id=dingding_conn_id,
                            message_type='markdown',
                            message={'title': title, 'text': content},
                            at_mobiles=at_mobiles,
                            at_all=at_all, ).send()

    return _dingding_sla_alert
